'''
Created on Jun 14, 2011

@author: Kathy Benitez
'''

from bitarray import bitarray
from random import choice
from operator import or_
from collections import defaultdict

class PolicyString(bitarray):
    """
    Represents a de-identification policy (non-hierarchical)
    """

    def __init__(self, initial=None):
        """
        Creates policy.  Possible initializers include array of booleans, string of binary
        """
        super(PolicyString, self).__init__(self, initial)
    def __repr__(self, verbose=False):
        if not verbose:
            return self.to01()
    def size(self):

        """
        Size or length of policy.
        >>> PolicyString(3).size()
        3L
        >>> PolicyString().size()
        0L
        """
        return self.length()    
    @classmethod
    def mostSpecific(PolicyString, length):
        """
        All ones (every vocabulary item in its own category)
        >>> PolicyString.mostSpecific(3)
        111
        """
        return PolicyString([1] * length)    
    @classmethod
    def mostGeneral(PolicyString, length):
        """
        All zeroes (every vocabulary item in one category per field--all items will be identical)
    
        >>> PolicyString.mostGeneral(3)
        000
    
        """
        return PolicyString([0] * length)    
    @classmethod
    def random(PolicyString, length):
        """
        Creates a random bitstring of the appropriate length.  This will
        be used in testing.  Ensure that this random policy is different from
        the most specific and most general cases.
        
        >>> PolicyString.random(5).size()
        5L
    
        """
        while True:
            policy = PolicyString([choice([0,1]) for x in range(length)]) 
            if (not policy.all() and policy.any()):
                return policy
    @staticmethod     
    def isGeneralization(general, specific):
        """
        General is a generalization of specific  iff all the places where specific has 
        a zero, general also has a zero
        
        Ancestor == More General
    
        >>> PolicyString.isGeneralization(PolicyString.mostGeneral(8), PolicyString.random(8))
        True
    
        >>> PolicyString.isGeneralization(PolicyString.mostSpecific(8), PolicyString.random(8))
        False
    
        >>> PolicyString.isGeneralization(PolicyString('1001001'), PolicyString('1001001'))
        False
    
        """
        if specific == general:
            return False
        return (~general | specific).all()     
    @staticmethod
    def isSpecification(specific, general):
        return PolicyString.isGeneralization(general, specific)
    @staticmethod
    def isGen(general,specific):
        return PolicyString.isGeneralization(general, specific)
    @staticmethod
    def isSpec(specific, general):
        return PolicyString.isSpecification(specific,general)
    @staticmethod
    def isAncestor(ancestor, descendant):
        """
        Deprecated
     
        Ancestor == More General
    
        """
        return PolicyString.isGeneralization(ancestor, descendant)
    def policyNeighbors(self, more=True, less=True):
        """
        Generates a list of policies which are neighbors to the
        original policy, that is, those that differ in one place.
        This function is used as the parent for both generalizations
        (less=True, more=False), and specifications (more=True, 
        less=False) 
        
        >>> sorted(PolicyString('101').policyNeighbors(more=True, less=True))
        [001, 100, 111]
    
        >>> PolicyString('101').policyNeighbors(more=False, less=False)
        []
    
        """
        if not more and not less:
            return []
        alternatives = []
        for position in range(self.size()):
            if less and self[position]:
                alternative = PolicyString(self)
                alternative[position] = 0
                alternatives.append(alternative)
            if more and not self[position]:
                alternative = PolicyString(self)
                alternative[position] = 1
                alternatives.append(alternative)
        return alternatives 
    def policyGeneralizations(self):
        """
        Generates a list of policies which are one level generalizations
        of the initial policy (i.e. have one position which was one in the
        original and is zero now
        
        >>> PolicyString('101').policyGeneralizations()
        [001, 100]
    
        >>> PolicyString('000').policyGeneralizations()
        []
    
        """
        return self.policyNeighbors(less=True, more=False)
    def policySpecifications(self):
        """
        Generates a list of policies which are one level specifications 
        of the initial policy (i.e. have one position which was zero in the
        original and is one now
    
        >>> PolicyString('101').policySpecifications()
        [111]
        
        >>> PolicyString('111').policySpecifications() 
        []
    
        """
        return self.policyNeighbors(less=False, more=True)
    def randomSpecification(self):
        """
        Choose one policy which is a one-level specification of the existing policy.
        It will have one division where the original grouped vocab items together.
        >>> PolicyString('101').randomSpecification()
        111
        >>> PolicyString('010').randomSpecification() in [bitarray('110'), bitarray('011')]
        True
        >>> PolicyString('111').randomSpecification()
        Traceback (most recent call last):
            . . .
        IndexError: policy already maximally specific
        """    
        try:
            child =  choice(self.policySpecifications())
            return child
        except IndexError:
            raise IndexError("policy already maximally specific")  
    def randomGeneralization(self):
        """
        >>> PolicyString('010').randomGeneralization()
        000
        >>> PolicyString('101').randomGeneralization() in [bitarray('001'), bitarray('100')]
        True
        >>> PolicyString('000').randomGeneralization()
        Traceback (most recent call last):
            ...
        IndexError: policy already maximally general
        """ 
        try:
            parent = choice(self.policyGeneralizations())
            return parent 
        except IndexError:
            raise IndexError("policy already maximally general")
    @staticmethod
    def levelsBetween(policyA, policyB):
        """
        >>> PolicyString.levelsBetween(bitarray('000'), bitarray('111'))
        3L
        >>> PolicyString.levelsBetween(bitarray('10000'), bitarray('10110'))
        2L
        >>> PolicyString.levelsBetween(bitarray('101'), bitarray('010'))
        Traceback (most recent call last):
            ...
        AssertionError: Policies not related
    
        """
        assert PolicyString.isGeneralization(policyA, policyB) or PolicyString.isGeneralization(policyB, policyA), "Policies not related"
        return(policyA ^ policyB).count()
    @staticmethod
    def halfpoint(generalLimit, specificLimit, bitweights):
        """
        >>> PolicyString.halfpoint(PolicyString('000'), PolicyString('110'), [1,1,1]) in [bitarray('100'), bitarray('010')]
        True
    
        If only one level apart, will return most specific policy, unchanged.
        >>> PolicyString.halfpoint(PolicyString('0010'), PolicyString('0110'), [1,1,1,1])
        0110
    
        >>> PolicyString.halfpoint(PolicyString('101'), PolicyString('010'), [1,1,1])
        Traceback (most recent call last):
            ...
        AssertionError: Policies not related
    
        >>> PolicyString.halfpoint(PolicyString('000'), PolicyString('110'), [1])
        Traceback (most recent call last):
            ...
        AssertionError: Not the appropriate number of bitweights
        """
        from util import weightedSample
    
        assert len(bitweights) == generalLimit.size(), "Not the appropriate number of bitweights"
        assert len(bitweights) == specificLimit.size(), "Not the appropriate number of bitweights"
        assert generalLimit.size()== specificLimit.size(), "Policies different sizes"
        assert PolicyString.isGeneralization(generalLimit, specificLimit), "Policies not related"
        diff = (~generalLimit) & specificLimit
        places = diff.search([True])
        mask = PolicyString.mostSpecific(generalLimit.size())
        positions = weightedSample(zip([bitweights[p] for p in places], places), len(places)/2)
        for p in positions:
            mask[p] = 0
        return specificLimit & mask 
    
class PolicyHandler():
    def __init__(self, vocab, fields):
        """
        Vocab takes the form of a dictionary:
        (fieldNo, eleNo) -> item
        i.e.
        (0, 0) -> "White alone"
        where 0 is Race and White is the first race on the list.
        """
        self.vocab = vocab
        self.fields = fields
        self.fieldLookup = dict([(v,k) for k,v in fields.iteritems()])
        self.vocabLists = defaultdict(dict)
        for w in vocab:
            self.vocabLists[w[0]][w[1]] = vocab[w]
#############
	self.divisions = {}
        for field in self.fields:
            numValues = len([i for i in vocab.keys() if i[0] == field])
            self.divisions[field] = [True] * (numValues-1)
###############
    def pprintPolicy(self, policy):
        #print self.vocabLists NOTE: MUST BE CALLED FROM PolicyConstructor, NOT PolicyHandler
        s = ""
	#changed assert to allow policy to be too big
        assert policy.size() >= len(self.vocab.keys()) - len(self.fields.keys())
        position = 0

        for field in self.fields:
            s += "%s:  %s"%(self.fields[field],self.vocabLists[field][0])
            for i in range(len(self.vocabLists[field]) - 1):
                if policy[position]:
                    s += " | %s"%self.vocabLists[field][i+1]
                else:
                    s += " %s"%self.vocabLists[field][i+1]
                position += 1
            s += "\n"
        position = 0
        for field in self.fields:
            for i in range(len(self.divisions[field])):
                self.divisions[field][i] = policy[position]
            position += 1
        print "Policy bitstring is %s"%(policy)
        print s



    def pprintPolicyReturn(self, policy,printZip=False,printAge=False):       #just like above method, but return instead of print, for Tri class
        #print self.vocabLists NOTE: MUST BE CALLED FROM PolicyConstructor, NOT PolicyHandler
        s = ""  
	#changed assert to allow policy to be too big
        assert policy.size() >= len(self.vocab.keys()) - len(self.fields.keys())
        
	if printZip:
		position=len(self.vocabLists[0])+len(self.vocabLists[1])+len(self.vocabLists[2])-3
		field=3
		s += "%s: \n \n(%s"%(self.fields[field],self.vocabLists[field][0])
            	for i in range(len(self.vocabLists[field]) - 1):
                	if policy[position]:
                 	   s += ") \n \n (%s"%self.vocabLists[field][i+1]    #note: replace hello with )),((
                	else:
                	   s += ", %s"%self.vocabLists[field][i+1]
               		position += 1
           	s += ")\n"
       		return s

	if printAge:
		position=len(self.vocabLists[0])+len(self.vocabLists[1])-2
		field=2
		s += "%s: \n \n(%s"%(self.fields[field],self.vocabLists[field][0])
            	for i in range(len(self.vocabLists[field]) - 1):
                	if policy[position]:
                 	   s += ") \n \n (%s"%self.vocabLists[field][i+1]    #note: replace hello with )),((
                	else:
                	   s += ", %s"%self.vocabLists[field][i+1]
               		position += 1
           	s += ")\n"
       		return s



	position=0
        for field in self.fields:
	  if not (self.fields[field]=='zip5' or self.fields[field]=='age'):
            s += "\n %s:  ((%s"%(self.fields[field],self.vocabLists[field][0])
            for i in range(len(self.vocabLists[field]) - 1):
                if policy[position]:
                    s += "))    (( %s"%self.vocabLists[field][i+1]
                else:
                    s += ", %s"%self.vocabLists[field][i+1]
                position += 1
            s += "))\n"
	   
        position = 0
        for field in self.fields:
            for i in range(len(self.divisions[field])):
                self.divisions[field][i] = policy[position]
            position += 1
       # print "Policy bitstring is %s"%(policy)
	
        return s

	




class PolicyConstructor(PolicyHandler):
    def __init__(self, vocab, fields):
        """
        Default policy is all separated
        """
        PolicyHandler.__init__(self, vocab, fields)
        self.divisions = {}
        for field in self.fields:
            numValues = len([i for i in vocab.keys() if i[0] == field])
            self.divisions[field] = [True] * (numValues-1)
    def getPolicy(self):
        return PolicyString(sum([self.divisions[field] for field in sorted(self.fields)],[]))
    def harvestPolicy(self, policy):
        assert policy.size() >= len(self.vocab.keys()) - len(self.fields.keys())
        position = 0
        for field in self.fields:
            for i in range(len(self.divisions[field])):
                self.divisions[field][i] = policy[position]
            position += 1
    def makeGeneral(self, field):
        self.divisions[field] = [False] * len(self.divisions[field])
    def makeSpecific(self, field):
        self.divisions[field] = [True] * len(self.divisions[field])
    def makeAllGeneral(self):
        for field in self.fields:
            self.makeGeneral(field)
    def makeAllSpecific(self):
        for field in self.fields:
            self.makeSpecific(field)
    def modNBins(self, field, n):
        field = self.fieldLookup[field]
        currentValue = self.vocabLists[field][0] / n
        for i in range(1, len(self.vocabLists[field])):
            newValue = self.vocabLists[field][i] / n
            self.divisions[field][i-1] = (newValue != currentValue)
            currentValue = newValue
    def matchNCharacters(self,field,n):
        field = self.fieldLookup[field]
        currentValue = str(self.vocabLists[field][0])[:n]
        for i in range(1, len(self.vocabLists[field])):
            newValue = str(self.vocabLists[field][i])[:n]
            self.divisions[field][i-1] = (newValue != currentValue)
            currentValue = newValue
class Dataset():
    def __init__(self, fields, dsDict, vocab=None):
        self.fields = fields 
        self.vocab = vocab
        self.ds = self.toAbstract(dsDict, vocab)
    def __getitem__(self,key):
        return self.ds[key]
    def applyPolicy(self,policy):
        newDs = {}
        translations = self.translationDictionary(policy)
	
        for group in self.ds:
            newGroup = tuple(sorted(sum([translations[word] for word in group], [])))
            newDs.setdefault(newGroup,0)
            newDs[newGroup] += self.ds[group]
        return newDs
#newDs looks just like abstract dict, accept with broad values
#for example: 
#((0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (2, 42), (2, 43), (2, 44), (2, 45):2)
    def translationDictionary(self,policy):
        newDict = {}
        fieldIndex = 0
        elementIndex = 0
        currentRange = [(fieldIndex, elementIndex)]
        for bit in policy:
            if bit:
                #there is a division, therefore, close the current range
                for word in currentRange:
                    newDict[word] = currentRange
                currentRange = [(fieldIndex, elementIndex+1)]
            else:
                currentRange.append((fieldIndex, elementIndex+1))
            elementIndex += 1
            if not (fieldIndex, elementIndex+1) in self.vocab:
                for word in currentRange:
                    newDict[word] = currentRange
                fieldIndex += 1
                elementIndex = 0    
                currentRange = [(fieldIndex, elementIndex)] 
        return newDict
    def getSpecifications(self,element):
        specifications = {}
        for e in element:
            field = self.fields[e[0]]
            specifications.setdefault(field,[])
            specifications[field].append(self.vocab[e])
        return specifications
    @staticmethod
    def toAbstract(ds, vocab):
        #invert the dictionary and create an abstract
        #representation of the dataset.  Represented as
        #a dictionary with keys a tuple of the values
        #in the record and value a count of the number
        #of times such a record appears
        if not vocab:
            raise TypeError("Must provide a valid vocabulary")
        minusV = dict([str(v),k] for k,v in vocab.iteritems())
        dset_A = {}
        try:
            for r in ds:
#changed v to r[v]
                newValue = tuple(sorted([minusV[r[v]] for v in r]))
                dset_A.setdefault(newValue,0)
                dset_A[newValue] += 1
            return dset_A
        except KeyError, k: 
            raise ValueError("%s not in vocabulary"%(k))
class Evaluator():
    def __init__(self, vocab, fields, dset, prior, ps):
        self.ds = Dataset(fields, dset, vocab)
        self.prior = prior
        self.previousPolicies = {}
        self.ps = ps
    
    def evaluatePolicy(self, policy):
	
        ds_policy = self.ds.applyPolicy(policy) 
        riskEstimate = 0
	
        for key in ds_policy:
            if key in self.prior:
                matches = self.prior[key]
            else:
                matches = self.ps.howManyMatches(self.ds.getSpecifications(key))
                self.prior[key] = matches
            if matches < ds_policy[key]:
                matches = ds_policy[key]
            riskEstimate += 1./matches * ds_policy[key]
        return riskEstimate

    def belowThreshold(self, newPolicy, threshold):
	#this function not needed, search.isPolicyAccepable is duplicate
        """
        policy is acceptable if a specification is acceptable.
        Policy is not acceptable if a generalization is not acceptable.
        Otherwise, must evaluate the policy itself
        
        """
        if newPolicy in self.previousPolicies:
            return self.previousPolicies <= threshold
        if reduce(or_, [PolicyString.isSpecification(p, newPolicy) and 
                        self.previousPolicies[p] <= threshold
                        for p in self.previousPolicies], False):
            return True
        if reduce(or_, [PolicyString.isGeneralization(p, newPolicy) and
                        self.previousPolicies[p] > threshold
                        for p in self.previousPolicies], False):
            return False
        else:
            """ evaluate the policy """
            value = self.evaluatePolicy(newPolicy) 
            self.nodesEvaluated += 1
            self.previousPolicies[newPolicy] = value
            return value <= self.cutoffValue 
